import open3d as o3d
import numpy as np
import time
import regiongrowing as reg

T1 = time.perf_counter()  # 计时开始
# ------------------------------读取点云---------------------------------------
pcd = o3d.io.read_point_cloud(r"D:\raw_data\point_cloud\01\no5.txt", format='xyz')
pcd = pcd.voxel_down_sample(voxel_size=0.01)  # 对原始点云进行降采样

# ------------------------------区域生长---------------------------------------
rg = reg.RegionGrow(pcd,  # 点云文件
                    min_pts_per_cluster=100,  # 每个聚类的最小点数
                    neighbour_number=30,  # 邻域搜索点数
                    theta_threshold=30,  # 平滑阈值（角度制）
                    curvature_threshold=0.05)  # 曲率阈值

# ---------------------------聚类结果分类保存----------------------------------
cluster_indices, curvature_array = rg.extract()  # 保存每个区域点云原始索引的双层list
print("筛选后聚类个数为", len(cluster_indices))
segment = []  # 存储分割结果的容器
pcd_index = 1
for ind in cluster_indices:
    # 遍历全部区域
    clusters_cloud = pcd.select_by_index(ind)
    r_color = np.random.uniform(0, 1, (1, 3))  # 分类点云随机赋色
    clusters_cloud.paint_uniform_color([r_color[:, 0], r_color[:, 1], r_color[:, 2]])
    segment.append(clusters_cloud)
    # # 保存到本地文件夹
    # pcd_path = r"./Region_Growing_pcd/cur_0.08/segment" + str(pcd_index) + ".pcd"
    # o3d.io.write_point_cloud(pcd_path, clusters_cloud)

    # 将点云数据保存到txt文件
    txt_path = r"./Region_Growing_txt2/cur_0.05/segment" + str(pcd_index) + ".txt"
    tmp_data = np.c_[np.array(clusters_cloud.points),
                     np.array(clusters_cloud.normals),
                     curvature_array[ind],
                     np.ones(len(ind)) * pcd_index]
    np.savetxt(txt_path, tmp_data)

    pcd_index += 1

T2 = time.perf_counter()  # 计时结束
print('点云预处理时间：', T2 - T1)

# -----------------------------结果可视化------------------------------------
o3d.visualization.draw_geometries(segment, window_name="区域生长分割",
                                  width=1024, height=768,
                                  left=50, top=50,
                                  mesh_show_back_face=False)
